package com.katze.ssh;

import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.session.ClientSession;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.time.Duration;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.Iterator;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

public class SshSessionPool implements AutoCloseable{
    private static final Logger log = LoggerFactory.getLogger(SshSessionPool.class);

    private final ConcurrentHashMap<String, SshSessionEntry> sessions = new ConcurrentHashMap<>();
    private final long idleTimeout; // 空闲超时时间（毫秒）
    private final Timer timer = new Timer();
    private final SshClient sshClient; // 共享的SSH客户端
    private final DefaultSessionListener sessionListener = new DefaultSessionListener();

    public SshSessionPool() {
        this(15 * 60 * 1000);
    }

    public SshSessionPool(long idleTimeout) {
        this.idleTimeout = idleTimeout;
        this.sshClient = SshClient.setUpDefaultClient();
        this.sshClient.addSessionListener(sessionListener);
        this.sshClient.start(); // 启动SSH客户端
        this.cleanup();
    }

    private void cleanup() {
        DateTimeFormatter formatter = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS");
        int minute = 5;
        LocalDateTime now = LocalDateTime.now();
        log.info("计算时间点:{}", now.format(formatter));

        // 间隔15分钟发送一次，但定时任务精确到毫秒，故即使程序启动时分钟是0，15，30，45也等待下一次
        LocalDateTime next = now.plusMinutes(minute - (now.getMinute() % minute));
        // 计算出首次执行时间点，分钟后去掉
        next = LocalDateTime.of(next.getYear(), next.getMonth(), next.getDayOfMonth(), next.getHour(), next.getMinute(), 0);

        long delay = Duration.between(now, next).toMillis();
        int period = minute * 60 * 1000;
        log.info("第一次执行时间点: {}, {}毫秒后开始, 间隔{}执行一次 ", next.format(formatter), delay, period);
        timer.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                int total = sessions.size();
                Iterator<Map.Entry<String, SshSessionEntry>> it = sessions.entrySet().iterator();
                while (it.hasNext()) {
                    SshSessionEntry se = it.next().getValue();
                    try {
                        if (se.isExpired(idleTimeout)) {
                            synchronized (se) {
                                if (se.isExpired(idleTimeout) && se.isOpen()) {
                                    se.close();
                                }
                                it.remove(); // 从池中移除
                            }
                        }
                    } catch (Exception e) {
                        log.error("{}巡检任务中操作失败", se.getAttribute(DefaultSessionListener.KEY), e);
                    }
                }
                log.debug("本次共清理{}个过期会话,当前剩余会话量: {}", total - sessions.size(), sessions.size());
            }
        }, delay, period);

    }

    public <T extends IChannel>T borrow(SshSessionKey key, Class<T> cs) throws Exception {
        return getClientSession(key).getChannel(cs);
    }

    /**
     * 获取客户端会话
     * @param key 会话标识
     */
    private SshSessionEntry getClientSession(SshSessionKey key) throws IOException {
        SshSessionEntry entry = sessions.get(key.getId());
        if (entry != null) {
            synchronized (entry) {
                if (entry.isOpen()) {
                    entry.updateLastAccess();
                    return entry;
                } else {
                    sessions.remove(key.getId(), entry); // 移除无效会话
                }
            }
        }
        ClientSession session = sshClient.connect(key.getUsername(), key.getHost(), key.getPort()).verify(key.getTimeout()).getSession();
        session.addPasswordIdentity(key.getPassword());
        session.setAttribute(DefaultSessionListener.KEY, key.getId());
        session.auth().verify(key.getTimeout());
        return sessions.computeIfAbsent(key.getId(), (k) -> new SshSessionEntry(session, key));
    }

    public SshSessionKey getSessionConfig(String key, Function<String, SecureShellSession> action) {
        SshSessionKey properties = sessionListener.sessionConfig(map -> map.get(key));
        if (properties == null) {
            SecureShellSession entity = action.apply(key);
            if (entity instanceof SshSessionKey prop) {
                properties = this.addSessionConfig(prop);
            } else {
                SshSessionKey prop = new SshSessionKey();
                prop.setId(entity.getId());
                prop.setUsername(entity.getUsername());
                prop.setHost(entity.getHost());
                prop.setPort(entity.getPort());
                prop.setPassword(entity.getPassword());
                prop.setTimeout(entity.getTimeout());
                properties = this.addSessionConfig(prop);
            }
        }
        return properties;
    }

    private SshSessionKey addSessionConfig(SshSessionKey properties) {
        return sessionListener.sessionConfig(map -> {
            if (!map.containsKey(properties.getId())) {
                map.put(properties.getId(), properties);
                log.info("Add session config {}:{} success!", properties.getId(), properties);
            } else {
                log.warn("{} is already exists: {}", properties.getId(), map.get(properties.getId()));
            }
            return properties;
        });
    }

    public boolean clear(String key) {
        SshSessionEntry entry = sessions.get(key);
        if (entry != null) {
            try {
                synchronized (entry) {
                    entry.close();
                }
                sessions.remove(key);
            } catch (Exception e) {
                log.error("{}关闭失败", key);
                return false;
            }
        }
        return true;
    }

    /**
     * 关闭所有会话以及巡检任务
     */
    @Override
    public void close() throws Exception {
        log.info("准备关闭SSH会话管理,当前会话量: {}", sessions.size());
        timer.cancel();
        // 关闭所有会话
        sessions.values().forEach(session -> {
            try {
                if (session.isOpen()) {
                    session.close();
                }
            } catch (Exception e) {
                log.error("{}关闭失败", session.getAttribute(DefaultSessionListener.KEY), e);
            }
        });
        sshClient.stop(); // 停止SSH客户端

        log.info("SSH会话管理已关闭.");
    }
}
